import numpy as np
from numpy import ndarray
from moudels.dct import DCT
from moudels.svd import SVD
class Watermarker:
    
    # 嵌入强度
    alpha:float = 0.2
    
    @classmethod
    def embed(cls, original_image:ndarray, watermark:ndarray):
        """水印嵌入函数"""
        dct_image = DCT.dct_transform(original_image)# 8*8 DCT
        watermarked_dct = dct_image.copy()
        
        #每个dct变换单位（8*8）的16个低频系数 64*64*4*4
        # 按照像素为单位
        for i in range(64):
            for j in range(64):
                dct_block:ndarray = dct_image[i,j]
                #提取低频系数并重构为64×64×4×4的矩阵
                low_freq = DCT.extract_low_freq(dct_block)

                #16个低频系数 4*4的矩阵进行SVD 取最大奇异值
                U, S, VH = SVD.svd(low_freq)
                # 嵌入水印 到最大奇异值
                S[0] += cls.alpha * watermark[i, j]   # 嵌入强度 * 水印
                # 嵌入后的低频系数
                low_freq = SVD.reconstruct(U, S, VH)
                          
                # 将嵌入水印后的低频系数重新放回DCT系数矩阵
                watermarked_dct[i,j] = DCT.reconstruct_dct_block(dct_block, low_freq)

        # idct 变换
        watermarked_image = DCT.idct_transform(watermarked_dct)        
        return watermarked_image

    @classmethod
    def extract(cls, watermarked_image:ndarray, original_image:ndarray):
        """水印提取函数"""
        original_dct = DCT.dct_transform(original_image)
        watermarked_dct = DCT.dct_transform(watermarked_image)
        
        #每个dct变换单位（8*8）的16个低频系数 64*64*4*4

        watermark_img = np.zeros((64,64),dtype=np.float64)
        
        # 按照像素为单位
        for i in range(64):
            for j in range(64):
                #16个低频系数 4*4的矩阵进行SVD 取最大奇异值
                original_svd_block = DCT.extract_low_freq(original_dct[i, j]).reshape(4,4)
                watermarked_svd_block =  DCT.extract_low_freq(watermarked_dct[i, j]).reshape(4,4)
                WU, WS, WVH = SVD.svd(watermarked_svd_block)
                U, S, VH = SVD.svd(original_svd_block)
                # 从最大奇异值 提取水印 
                watermark_img[i,j] = (WS[0] - S[0]) / cls.alpha   # 嵌入结果 /嵌入强度 = 水印

        return watermark_img
